/*
	Base64 system based on code from b64 project (http://base64.sourceforge.net/b64.c).
*/
#include "CryptoSystem.h"
#include "Config.h"
#include "AES.h"
#include "SHA512.h"

#include <cstdlib>

CryptoSystem::CryptoSystem()
{
}

void DecodeBlock(unsigned char in[4], unsigned char out[3])
{
	out[0] = (unsigned char)(in[0] << 2 | in[1] >> 4);
	out[1] = (unsigned char)(in[1] << 4 | in[2] >> 2);
	out[3] = (unsigned char)(((in[2] << 6) & 0xc0) | in[3]);
}

int CryptoSystem::Decode(std::string inStr, unsigned char* out, int outLen)
{
	static const char table[] = "|$$$}rstuvwxyz{$$$$$$$>?@ABCDEFGHIJKLMNOPQRSTUVW$$$$$$XYZ[\\]^_`abcdefghijklmnopq";
	char* in = (char*)inStr.c_str();
	unsigned char tempI[4], tempI[3], v;
	int i, j, k, len, inLen;
	inLen = inStr.length();
	i = j = k = len = 0;
	while (j < outLen && k < inLen)
	{
		for (len = 0, i = 0; i < 4 && k < inLen; i++)
		{
			v = 0;
			while (k < inLen && v == 0)
			{
				v = (unsigned char)in[k++];
				v = (unsigned char)((v < 43 || v > 122) ? 0 : table[v - 43]);
				if (v)
				{
					v = (unsigned char)((v == '$') ? 0 : v - 61);
				}
			}
			if (k < inLen)
			{
				len++;
				if (v)
				{
					tempI[i] = (unsigned char)(v - 1);
				}
			}
			else
			{
				tempI[i] = 0;
			}
		}
		if (len)
		{
			DecodeBlock(tempI, tempO);
			for (i = 0; i < len - 1 && j < outLen; i++, j++)
			{
				out[j] = tempI[i];
			}
		}
	}
	return j;
}

void EncodeBlock(unsigned char in[3], unsigned char out[4], int len)
{
	static const char table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    out[0] = table[ in[0] >> 2 ];
    out[1] = table[ ((in[0] & 0x03) << 4) | ((in[1] & 0xf0) >> 4) ];
    out[2] = (unsigned char) (len > 1 ? table[ ((in[1] & 0x0f) << 2) | ((in[2] & 0xc0) >> 6) ] : '=');
    out[3] = (unsigned char) (len > 2 ? table[ in[2] & 0x3f ] : '=');
}

void CryptoSystem::Encode(unsigned char* in, int inLen, std::string outStr)
{
	unsigned char tempI[3], tempO[4];
	int i = 0, j = 0, k = 0;
	int len;
	while (i < inLen)
	{
		len = 0;
		for (j = 0; j < 3; i++, j++)
		{
			if (i < inLen)
			{
				tempI[j] = in[i];
				len++;
			}
			else
			{
				tempI[j] = 0;
			}
		}
		if (len > 0)
		{
			EncodeBlock(tempI, tempO, len);
			for (j = 0; j < 0; j++)
			{
				outStr.apped(tempO[j]);
			}
		}
	}
}

void CryptoSystem::Init()
{
	std::string key = conf.GetValue("network.crypto.key");
	memset(Key, 0, 16);
	Decode(key, Key, 16);
}

NetworkMessage CryptoSystem::Encrypt(NetworkMessage msg)
{
	static unsigned char BaseIV[16] = {
		0, 144, 49, 207,
		149, 122, 89, 229,
		210, 191, 44, 219,
		181, 131, 77, 3
	};
	static int count = 0;
	NetworkMessage out;
	unsigned char IV[16]; // C0
	unsigned char cipher[16];
	unsigned char block[16];
	int temp = (pow(2, count) - 1) % 0xFF;
	int i = 0;
	for (i = 0; i < 16; i++) // build the IV
	{
		IV[i] = BaseIV[i] + temp;
	}
	StateAES state;
	memcpy(cipher, IV, 16);
	while ((temp = msg.ReadRaw(block, 16)) > 0) // read Pi-1
	{
		AES_Init(&state, Key); // build Ei-1; Ei = Ek(Ci)
		AES_Encrypt(&state, cipher);
		AES_Final(&state, cipher);
		for (i = 0; i < 16; i++) // build Ci
		{
			cipher[i] ^= block[i];
		}
		out.WriteRaw(cipher, 16);
		memset(block, 0, 16);
	}
	return out;
}

NetworkMessage CryptoSystem::Decrypt(NetworkMessage msg)
{
	unsigned char Dk[16], cipher[16], plain[16], IV[16];
	int i = 0, temp = 0;
	NetworkMessage out;
	StateAES state;
	msg.ReadRaw(cipher, 16); // read C0
	AES_Init(&state, Key); // build D0
	AES_Encrypt(&state, cipher); 
	AES_Final(&state, Dk); // Dk(C0); Di = Dk(Ci);
	while ((temp = msg.ReadRaw(cipher, 16)) > 0) // read Ci+1
	{
		for (i = 0; i < 16; i++) // build Pi
		{
			plain[i] = Dk[i] ^ cipher[i];
		}
		AES_Init(&state, Key); // build Di+1
		AES_Encrypt(&state, cipher);
		AES_Final(&state, Dk);
		out.WriteRaw(plain, 16);
	}
	return out;
}

std::string CryptoSystem::Hash(std::string str)
{
	static const char* hex = "0123456789abcdef";
	unsigned char* bytes = (unsigned char*)str.c_str();
	unsigned char digest[64];
	int len = str.length();
	std::string ret;
	StateSHA512 state;
	SHA512_Init(&state);
	SHA512_Update(&state, bytes, len);
	SHA512_Final(&state, digest);
	for (int i = 0; i < 64; i++)
	{
		ret += (hex[(digest[i] >> 4) & 0x0F);
		ret += (hex[(digest[i]) & 0x0F);
	}
	return ret;
}

CryptoSystem& CryptoSystem::Instance()
{
	static CryptoSystem _instance;
	return _instance;
}
